package com.example.shardingjdbc.config;

import com.alibaba.druid.util.StringUtils;
import com.google.common.collect.Lists;
import com.google.common.collect.Range;
import lombok.extern.slf4j.Slf4j;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingAlgorithm;
import org.apache.shardingsphere.api.sharding.complex.ComplexKeysShardingValue;

import java.text.SimpleDateFormat;
import java.util.*;

@Slf4j
public class MonthBaseTablePreciseComplexKeysShardingAlgorithm implements ComplexKeysShardingAlgorithm<Date> {

    @Override
    public Collection<String> doSharding(Collection availableTargetNames, ComplexKeysShardingValue shardingValue) {
        System.out.println("222222222222222222222222");
        Object o = shardingValue.getColumnNameAndRangeValuesMap().get("data_time");

        Range<?> range = (Range<?>) o;

        // 获取范围的起始日期和结束日期
        Object lowerEndpoint = range.lowerEndpoint();
        Object upperEndpoint = range.upperEndpoint();

        // 确保它们是 Date 类型（这里假设它们确实是 Date 类型）
        Date startDate = (Date) lowerEndpoint;
        Date endDate = (Date) upperEndpoint;

        // 处理日期范围
        Set<String> result = new HashSet<>();
        SimpleDateFormat sdf = new SimpleDateFormat("MM");

        // 遍历从 startDate 到 endDate 的每个月份
        Calendar calendar = Calendar.getInstance();
        calendar.setTime(startDate);

        while (!calendar.getTime().after(endDate)) {
            String month = sdf.format(calendar.getTime());
            String tableSuffix = String.format("%02d", Integer.parseInt(month));
            for (Object table : availableTargetNames) {
                String tableStr = (String) table;
                if (tableStr.endsWith("_" + tableSuffix)) {
                    result.add(tableStr);
                }
            }
            // 移动到下一个月
            calendar.add(Calendar.MONTH, 1);
        }

        return result.isEmpty() ? availableTargetNames : result;
    }

    /**
     * 根据日期获取数据表名的范围
     *
     * @param dateFormat
     * @param startTime
     * @param endTime
     * @return
     */
    private static List<String> getRangeNameList(SimpleDateFormat dateFormat, String range, Date startTime, Date endTime) {

        List<String> result = Lists.newArrayList();
        // 定义日期实例
        Calendar dd = Calendar.getInstance();
        dd.setTime(ShardingDateUtils.beginOfDate(startTime,range));

        while (dd.getTime().before(endTime) || dd.getTime().equals(endTime)) {
            result.add(dateFormat.format(dd.getTime()));
            // 根据策略进行 + 1
            dd.add(ShardingDateUtils.getCalendarDateRange(range), 1);
        }
        return result;
    }

    /**
     * 分片拿取值
     *
     * @param columnNameAndShardingValuesMap
     * @param key
     * @return
     */
    protected Collection getShardingValue(Map<String, Collection> columnNameAndShardingValuesMap, String key) {
        Collection valueSet = new ArrayList();

        if (StringUtils.isEmpty(key)) {
            return valueSet;
        }

        if (columnNameAndShardingValuesMap.containsKey(key.toLowerCase())) {
            key = key.toLowerCase();
        } else if (columnNameAndShardingValuesMap.containsKey(key.toUpperCase())) {
            key = key.toUpperCase();
        } else {
            return valueSet;
        }
        if (columnNameAndShardingValuesMap.get(key) instanceof Collection) {
            valueSet.addAll(columnNameAndShardingValuesMap.get(key));
        } else {
            valueSet.add(columnNameAndShardingValuesMap.get(key));
        }
        return valueSet;
    }
}